#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"

#include "crypto/Crypto.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir


using namespace mlir;
using namespace vector;
using namespace mlir::arith;

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {

class PadMessageLowering : public OpRewritePattern<crypto::PadMessage_OP> {
public:
  using OpRewritePattern<crypto::PadMessage_OP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::PadMessage_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    Value message = op->getOperand(0);
    Value message_len = op->getOperand(1);

    // constants
    Value sixtyfour_I32 = rewriter.create<ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(64));
    Value eight_I32 = rewriter.create<ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(8));
    Value one_I32 = rewriter.create<ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
    Value one_I8 = rewriter.create<ConstantOp>(loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(0x80));
    Value zero_I8 = rewriter.create<ConstantOp>(loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(0));
    Value c8_I64 = rewriter.create<ConstantOp>(loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(8));
    Value c1 = rewriter.create<ConstantIndexOp>(loc, /*value=*/1);
    Value c8 = rewriter.create<ConstantIndexOp>(loc, /*value=*/8);
    // types used
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());

    // calculate the padding size
    Value rem = rewriter.create<arith::RemSIOp>(loc, /*lhs=*/message_len, /*rhs=*/sixtyfour_I32);
    Value diff = rewriter.create<arith::SubIOp>(loc, /*lhs=*/sixtyfour_I32, /*rhs=*/rem);
    Value isPadTwo = rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::sle, /*lhs=*/diff, /*rhs=*/eight_I32);
    Value padSize1 = rewriter.create<arith::AddIOp>(loc, /*lhs=*/message_len, /*rhs=*/diff);
    Value padSize2 = rewriter.create<arith::AddIOp>(loc, /*lhs=*/padSize1, /*rhs=*/sixtyfour_I32);
    Value padSize = rewriter.create<arith::SelectOp>(loc, /*condition=*/isPadTwo, /*true_value=*/padSize2, /*false_value=*/padSize1);
    Value padSizeIdx = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/padSize);

    // allocate memory for padded message
    Value padded_message = rewriter.create<memref::AllocOp>(loc, /*memrefType=*/dynI8MemrefType, /*dynamicSizes=*/ValueRange{padSizeIdx});
    // auto msgDevAllocOp = rewriter.create<gpu::AllocOp>(
    //       loc, 
    //       /*memref=*/dynI8MemrefType, 
    //       /*asyncDependencies=*/ValueRange{}, 
    //       /*dynamicSizes=*/ValueRange{padSizeIdx}, 
    //       /*symbolOperands=*/ValueRange{}
    //     );
    // Value padded_message = msgDevAllocOp.getResult(0);

    // copy the original message to the padded message
    rewriter.create<memref::CopyOp>(loc, /*source=*/message, /*target=*/padded_message);

    // fill the last byte with 0x80
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/message_len);
    rewriter.create<memref::StoreOp>(loc, /*value=*/one_I8, /*memref=*/padded_message, /*indices=*/ValueRange{msgLenIdx});
    Value msgLenPlusOneIdx = rewriter.create<arith::AddIOp>(loc, /*lhs=*/msgLenIdx, /*rhs=*/c1);
    Value lenPos =  rewriter.create<arith::SubIOp>(loc, /*lhs=*/padSizeIdx, /*rhs=*/c8); 
    // fill the padded part with zeros
    rewriter.create<scf::ForOp>(
      loc, /*lowerBound=*/msgLenPlusOneIdx, /*upperBound=*/lenPos, /*step=*/c1, /*iterArgs=*/ValueRange{},
      /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
        rewriter.create<memref::StoreOp>(loc, /*value=*/zero_I8, /*memref=*/padded_message, /*indices=*/ValueRange{ivs});
        rewriter.create<scf::YieldOp>(loc);
      });

    // calculate the message length in bits and store it in the padded message
    Value extLen = rewriter.create<arith::ExtUIOp>(loc, /*out=*/rewriter.getI64Type(), /*in=*/message_len);
    Value msgBits = rewriter.create<arith::MulIOp>(loc, /*lhs=*/extLen, /*rhs=*/c8_I64);
    Value t0 = rewriter.create<vector::BroadcastOp>(loc, /*vector=*/VectorType::get({1}, rewriter.getI64Type()), /*source=*/msgBits);
    Value t1 = rewriter.create<vector::BitCastOp>(loc, /*result=*/VectorType::get({8}, rewriter.getI8Type()), /*source=*/t0);
    SmallVector<int64_t> mask = {7, 6, 5, 4, 3, 2, 1, 0};
    Value msgBitsVec = rewriter.create<vector::ShuffleOp>(loc, /*v1=*/t1, /*v2=*/t1, /*odsArg2=*/mask);
    rewriter.create<vector::StoreOp>(loc, /*valueToStore=*/msgBitsVec, /*base=*/padded_message, /*indices=*/lenPos);
    
    rewriter.replaceOp(op, {padded_message});
    return success();
  }
};

class PadMessageGetPadmemLowering : public OpRewritePattern<crypto::PadMessageGetPadmem_OP> {
public:
  using OpRewritePattern<crypto::PadMessageGetPadmem_OP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::PadMessageGetPadmem_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    Value message = op->getOperand(0);
    Value message_len = op->getOperand(1);
    Value message_num = op->getOperand(2);

    // cast to index type for iteration
    Value cn = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/message_len);
    Value cN = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/message_num);

    // constants
    Value c0 = rewriter.create<ConstantIndexOp>(loc, /*value=*/0);
    Value c1 = rewriter.create<ConstantIndexOp>(loc, /*value=*/1);
    Value c64I32 = rewriter.create<ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(64));
    Value c8I32 = rewriter.create<ConstantOp>(loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(8));
    // types used
    auto dynI8MemrefType = MemRefType::get({mlir::ShapedType::kDynamic}, rewriter.getI8Type());

    // calculate the padding size
    Value rem = rewriter.create<arith::RemSIOp>(loc, /*lhs=*/message_len, /*rhs=*/c64I32);
    Value diff = rewriter.create<arith::SubIOp>(loc, /*lhs=*/c64I32, /*rhs=*/rem);
    Value isPadTwo = rewriter.create<arith::CmpIOp>(loc, /*predicate=*/arith::CmpIPredicate::sle, /*lhs=*/diff, /*rhs=*/c8I32);
    Value padSize1 = rewriter.create<arith::AddIOp>(loc, /*lhs=*/message_len, /*rhs=*/diff);
    Value padSize2 = rewriter.create<arith::AddIOp>(loc, /*lhs=*/padSize1, /*rhs=*/c64I32);
    Value padSize = rewriter.create<arith::SelectOp>(loc, /*condition=*/isPadTwo, /*true_value=*/padSize2, /*false_value=*/padSize1);
    Value padSizeIdx = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/padSize);
    Value totalSizeIdx = rewriter.create<arith::MulIOp>(loc, /*lhs=*/padSizeIdx, /*rhs=*/cN);

    // allocate total memory for padded message
    Value padded_message = rewriter.create<memref::AllocOp>(loc, /*memrefType=*/dynI8MemrefType, /*dynamicSizes=*/ValueRange{totalSizeIdx});

    // copy the original message to the padded message
    rewriter.create<scf::ForOp>(
      loc, /*lowerBound=*/c0, /*upperBound=*/cN, /*step=*/c1, /*iterArgs=*/ValueRange{},
      /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
        Value srcBase = rewriter.create<arith::MulIOp>(loc, /*lhs=*/cn, /*rhs=*/ivs);
        Value dstBase = rewriter.create<arith::MulIOp>(loc, /*lhs=*/padSizeIdx, /*rhs=*/ivs);
        rewriter.create<scf::ForOp>(
          loc, /*lowerBound=*/c0, /*upperBound=*/cn, /*step=*/c1, /*iterArgs=*/ValueRange{},
          /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value offset, ValueRange iargs){
            Value srcIndice = rewriter.create<arith::AddIOp>(loc, /*lhs=*/srcBase, /*rhs=*/offset);
            Value dstIndice = rewriter.create<arith::AddIOp>(loc, /*lhs=*/dstBase, /*rhs=*/offset);
            Value ele = rewriter.create<memref::LoadOp>(loc, /*memref=*/message, /*indices=*/ValueRange{srcIndice});
            rewriter.create<memref::StoreOp>(loc, /*value=*/ele, /*memref=*/padded_message, /*indices=*/ValueRange{dstIndice});
            rewriter.create<scf::YieldOp>(loc);
          });
        rewriter.create<scf::YieldOp>(loc);
      });
    
    rewriter.replaceOp(op, {padded_message, padSizeIdx});
    return success();
  }
};

class PadMessagePadoneLowering : public OpRewritePattern<crypto::PadMessagePadone_OP> {
public:
  using OpRewritePattern<crypto::PadMessagePadone_OP>::OpRewritePattern;

  LogicalResult matchAndRewrite(crypto::PadMessagePadone_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op->getLoc();
    Value pad_mem = op->getOperand(0);
    Value message_len = op->getOperand(1);

    // constants
    Value one_I8 = rewriter.create<ConstantOp>(loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(0x80));
    Value c0I8 = rewriter.create<ConstantOp>(loc, rewriter.getI8Type(), rewriter.getI8IntegerAttr(0));
    Value c1 = rewriter.create<ConstantIndexOp>(loc, /*value=*/1);
    Value c8 = rewriter.create<ConstantIndexOp>(loc, /*value=*/8);
    Value c8I64 = rewriter.create<ConstantOp>(loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(8));
    // variables used
    Value msgLenIdx = rewriter.create<arith::IndexCastOp>(loc, /*out=*/rewriter.getIndexType(), /*in=*/message_len);
    Value padSizeIdx = rewriter.create<memref::DimOp>(loc, /*memref=*/pad_mem, /*index=*/0);

    // fill the last byte with 0x80
    rewriter.create<memref::StoreOp>(loc, /*value=*/one_I8, /*memref=*/pad_mem, /*indices=*/ValueRange{msgLenIdx});

    // fill the padded part with zeros
    rewriter.create<scf::ForOp>(
      loc, /*lowerBound=*/rewriter.create<arith::AddIOp>(loc, /*lhs=*/msgLenIdx, /*rhs=*/c1), 
      /*upperBound=*/rewriter.create<arith::SubIOp>(loc, /*lhs=*/padSizeIdx, /*rhs=*/c8), 
      /*step=*/c1, /*iterArgs=*/ValueRange{},
      /*odsArg4=*/[&](OpBuilder &builder, Location loc, Value ivs, ValueRange iargs){
        rewriter.create<memref::StoreOp>(loc, /*value=*/c0I8, /*memref=*/pad_mem, /*indices=*/ValueRange{ivs});
        rewriter.create<scf::YieldOp>(loc);
      });

    // calculate the message length in bits and store it in the padded message
    Value extLen = rewriter.create<arith::ExtUIOp>(loc, /*out=*/rewriter.getI64Type(), /*in=*/message_len);
    Value msgBits = rewriter.create<arith::MulIOp>(loc, /*lhs=*/extLen, /*rhs=*/c8I64);
    Value t0 = rewriter.create<vector::BroadcastOp>(loc, /*vector=*/VectorType::get({1}, rewriter.getI64Type()), /*source=*/msgBits);
    Value t1 = rewriter.create<vector::BitCastOp>(loc, /*result=*/VectorType::get({8}, rewriter.getI8Type()), /*source=*/t0);
    SmallVector<int64_t> mask = {7, 6, 5, 4, 3, 2, 1, 0};
    Value lenPos =  rewriter.create<arith::SubIOp>(loc, /*lhs=*/padSizeIdx, /*rhs=*/c8);
    Value msgBitsVec = rewriter.create<vector::ShuffleOp>(loc, /*v1=*/t1, /*v2=*/t1, /*odsArg2=*/mask);
    rewriter.create<vector::StoreOp>(loc, /*valueToStore=*/msgBitsVec, /*base=*/pad_mem, /*indices=*/lenPos);

    rewriter.eraseOp(op);
    return success();
  }
};

} // namespace

void populateLowerPadMessagePatterns(RewritePatternSet &patterns) {
  // patterns.add<PadMessageLowering>(patterns.getContext());
  patterns.add<PadMessageGetPadmemLowering>(patterns.getContext());
  patterns.add<PadMessagePadoneLowering>(patterns.getContext());
}

namespace {
class LowerPadMessagePass : public PassWrapper<LowerPadMessagePass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerPadMessagePass)
  LowerPadMessagePass() = default;
  LowerPadMessagePass(const LowerPadMessagePass &) {}

  StringRef getArgument() const final { return "lower-crypto-padmessage"; }
  StringRef getDescription() const final { return "Lower Crypto PadMessage Op."; }

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<func::FuncDialect,
                mlir::LLVM::LLVMDialect,
                memref::MemRefDialect,
                arith::ArithDialect,
                vector::VectorDialect,
                affine::AffineDialect,
                gpu::GPUDialect,
                scf::SCFDialect
                >();
  }
};
} // end anonymous namespace.

void LowerPadMessagePass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<
  func::FuncDialect,
  mlir::LLVM::LLVMDialect,
  memref::MemRefDialect,
  arith::ArithDialect,
  affine::AffineDialect,
  gpu::GPUDialect,
  vector::VectorDialect,
  scf::SCFDialect
  >();
  target.addLegalOp<
  ModuleOp,
  func::FuncOp,
  func::ReturnOp
  >();

  RewritePatternSet patterns(context);
  populateLowerPadMessagePatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerPadMessagePass() { PassRegistration<LowerPadMessagePass>(); }

std::unique_ptr<Pass> createLowerPadMessagePass(){ return std::make_unique<LowerPadMessagePass>(); }
